package com.mahiru.mahiruaiagent.tools;

import cn.hutool.core.io.FileUtil;
import cn.hutool.http.HttpRequest;
import cn.hutool.http.HttpResponse;
import cn.hutool.http.HttpUtil;
import com.mahiru.mahiruaiagent.constant.FileConstant;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.tool.annotation.Tool;
import org.springframework.ai.tool.annotation.ToolParam;

import java.io.File;
import java.util.Arrays;
import java.util.HashSet;
import java.util.Set;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * 从指定url中下载资源
 *
 * @Author Enkidu
 * @Date 2025/7/8 0:15
 */
public class ResourceDownloadTool {

    private static final Logger logger = LoggerFactory.getLogger(ResourceDownloadTool.class);

    // 支持的图片MIME类型
    private static final Set<String> SUPPORTED_IMAGE_MIME_TYPES = new HashSet<>(Arrays.asList(
            "image/jpeg", "image/png", "image/gif", "image/webp", "image/bmp", "image/tiff"));

    // 图片URL模式匹配
    private static final Pattern IMAGE_URL_PATTERN = Pattern.compile(".*\\.(jpg|jpeg|png|gif|webp|bmp|tiff)$",
            Pattern.CASE_INSENSITIVE);

    @Tool(description = "Download a resource from a given URL")
    public String downloadResource(@ToolParam(description = "URL of the resource to download") String url,
            @ToolParam(description = "Name of the file to save the downloaded resource") String fileName) {
        String fileDir = FileConstant.FILE_SAVE_DIR + "/download";
        String filePath = fileDir + "/" + fileName;
        try {
            // 创建目录
            FileUtil.mkdir(fileDir);

            // 检查URL是否直接是图片URL
            boolean isImageUrl = IMAGE_URL_PATTERN.matcher(url).matches();

            // 如果URL看起来不像图片，先检查内容类型
            if (!isImageUrl) {
                try (HttpResponse response = HttpRequest.head(url).execute()) {
                    String contentType = response.header("Content-Type");
                    logger.info("URL: {}, Content-Type: {}", url, contentType);

                    if (contentType == null || !SUPPORTED_IMAGE_MIME_TYPES.contains(contentType.toLowerCase())) {
                        // 如果不是图片类型，尝试从页面中提取图片URL
                        String imageUrl = extractImageUrlFromWebPage(url);
                        if (imageUrl != null) {
                            logger.info("提取到图片URL: {}", imageUrl);
                            url = imageUrl;
                        } else {
                            return "错误：提供的URL不是图片资源，也无法从中提取图片: " + url;
                        }
                    }
                }
            }

            // 使用 Hutool 的 downloadFile 方法下载资源
            HttpUtil.downloadFile(url, new File(filePath));
            return "资源下载成功，保存路径: " + filePath;
        } catch (Exception e) {
            logger.error("下载资源时出错", e);
            return "下载资源时出错: " + e.getMessage();
        }
    }

    /**
     * 专门用于下载图片的方法
     */
    @Tool(description = "Download an image from a given URL, ensuring the resource is actually an image")
    public String downloadImage(@ToolParam(description = "URL of the image to download") String url,
            @ToolParam(description = "Name of the file to save the downloaded image") String fileName) {
        String fileDir = FileConstant.FILE_SAVE_DIR + "/images";
        String filePath = fileDir + "/" + fileName;

        // 确保文件名有正确的扩展名
        if (!hasImageExtension(fileName)) {
            fileName = addImageExtension(fileName, url);
            filePath = fileDir + "/" + fileName;
        }

        try {
            // 创建目录
            FileUtil.mkdir(fileDir);

            // 检查URL是否指向图片
            try (HttpResponse response = HttpRequest.head(url).execute()) {
                String contentType = response.header("Content-Type");
                logger.info("URL: {}, Content-Type: {}", url, contentType);

                if (contentType == null || !SUPPORTED_IMAGE_MIME_TYPES.contains(contentType.toLowerCase())) {
                    // 如果不是图片类型，尝试从页面中提取图片URL
                    String imageUrl = extractImageUrlFromWebPage(url);
                    if (imageUrl != null) {
                        logger.info("提取到图片URL: {}", imageUrl);
                        url = imageUrl;
                    } else {
                        return "错误：提供的URL不是图片资源，也无法从中提取图片: " + url;
                    }
                }
            }

            // 下载图片
            HttpUtil.downloadFile(url, new File(filePath));
            return "图片下载成功，保存路径: " + filePath;
        } catch (Exception e) {
            logger.error("下载图片时出错", e);
            return "下载图片时出错: " + e.getMessage();
        }
    }

    /**
     * 从网页中提取图片URL
     */
    private String extractImageUrlFromWebPage(String url) {
        try {
            String html = HttpUtil.get(url);

            // 使用正则表达式提取图片URL
            Pattern pattern = Pattern.compile("<img[^>]+src=[\"']([^\"']+)[\"']", Pattern.CASE_INSENSITIVE);
            Matcher matcher = pattern.matcher(html);

            // 寻找第一个有效的图片URL
            while (matcher.find()) {
                String imgUrl = matcher.group(1);

                // 处理相对URL
                if (imgUrl.startsWith("/")) {
                    // 提取域名
                    Pattern domainPattern = Pattern.compile("(https?://[^/]+)");
                    Matcher domainMatcher = domainPattern.matcher(url);
                    if (domainMatcher.find()) {
                        String domain = domainMatcher.group(1);
                        imgUrl = domain + imgUrl;
                    }
                } else if (!imgUrl.startsWith("http")) {
                    // 相对路径但不是以/开头
                    String baseUrl = url.substring(0, url.lastIndexOf('/') + 1);
                    imgUrl = baseUrl + imgUrl;
                }

                // 检查是否是图片URL
                if (IMAGE_URL_PATTERN.matcher(imgUrl).matches()) {
                    return imgUrl;
                }

                // 检查内容类型
                try (HttpResponse response = HttpRequest.head(imgUrl).execute()) {
                    String contentType = response.header("Content-Type");
                    if (contentType != null && SUPPORTED_IMAGE_MIME_TYPES.contains(contentType.toLowerCase())) {
                        return imgUrl;
                    }
                } catch (Exception e) {
                    // 忽略单个图片检查错误，继续下一个
                    logger.warn("检查图片URL时出错: {}", e.getMessage());
                }
            }

            // 没有找到有效图片
            return null;
        } catch (Exception e) {
            logger.error("从网页提取图片时出错", e);
            return null;
        }
    }

    /**
     * 检查文件名是否有图片扩展名
     */
    private boolean hasImageExtension(String fileName) {
        return IMAGE_URL_PATTERN.matcher(fileName).matches();
    }

    /**
     * 为文件名添加适当的图片扩展名
     */
    private String addImageExtension(String fileName, String url) {
        // 从URL中提取扩展名
        Matcher matcher = IMAGE_URL_PATTERN.matcher(url);
        if (matcher.find()) {
            String extension = matcher.group(1);
            return fileName + "." + extension;
        }

        // 默认使用jpg扩展名
        return fileName + ".jpg";
    }
}
